// hisi_l0_mem_pool_core.c
#include <linux/sort.h>
#include <linux/slab.h>
#include <linux/mm.h>
#include <linux/module.h>
#include <linux/cpu.h>
#include "hisi_l0_mem_pool.h"
#include "hisi_l3t.h"

// Move global variables from main file
int base_node;
cache_tuner l0_cache_tuner __read_mostly;
struct kobject *dump_l0_pool_kobj;

// Module parameters
#define DEFAULT_MAX_NUMA_CAPACITY 20
uint max_numa_capacity = DEFAULT_MAX_NUMA_CAPACITY;
module_param(max_numa_capacity, uint, 0644);
MODULE_PARM_DESC(max_numa_capacity, "Maximum capacity per NUMA node in MB");

uint pre_alloc_size = 0;
module_param(pre_alloc_size, uint, 0644);
MODULE_PARM_DESC(pre_alloc_size, "Pre-allocation size per NUMA node in MB");

uint preinit_cache_levels = 0x00;
module_param(preinit_cache_levels, uint, 0644);
MODULE_PARM_DESC(preinit_cache_levels, "Bitmask of cache levels to preinitialize");

bool prealloc_node_enabled[MAX_NUMA_NODES] = {false};
bool prealloc_nodes_parsed = false;
char *prealloc_numa_nodes = "";
module_param(prealloc_numa_nodes, charp, 0644);
MODULE_PARM_DESC(prealloc_numa_nodes, "Comma-separated list of NUMA nodes to pre-allocate (e.g., '0,1,2')");

// - early_numa_node_count_init
static int early_numa_node_count_init(void)
{
    int node_count = num_online_nodes();
    
    if (node_count < 1) {
        pr_err("No NUMA nodes found\n");
        return -ENODEV;
    }
    
    atomic_set(&l0_cache_tuner.numa_node_count, node_count);
    return 0;
}

static int __init parse_numa_token(const char *token, int *node_out)
{
    int node;
    const char *p;

    // Check if token is empty
    if (!token || *token == '\0') {
        pr_err("Empty token found in numa node list\n");
        return -EINVAL;
    }

    // Check if token contains only digits
    for (p = token; *p != '\0'; p++) {
        if (*p < '0' || *p > '9') {
            pr_err("Invalid character '%c' in token '%s'\n", *p, token);
            return -EINVAL;
        }
    }

    // Convert to integer
    if (kstrtoint(token, 10, &node)) {
        pr_err("Failed to convert token '%s' to integer\n", token);
        return -EINVAL;
    }

    *node_out = node;
    return 0;
}

static int __init update_global_state(const bool *valid_nodes, int node_count)
{
    int i;

    if (!valid_nodes) {
        pr_err("Invalid valid_nodes array\n");
        return -EINVAL;
    }

    for (i = 0; i < node_count && i < MAX_NUMA_NODES; i++) {
        prealloc_node_enabled[i] = valid_nodes[i];
    }

    prealloc_nodes_parsed = true;
    return 0;
}

static int __init parse_and_validate_prealloc_numa_nodes(void)
{
    char *buffer = NULL, *token, *tmp;
    bool valid_nodes[MAX_NUMA_NODES] = {false};
    int node_count = get_numa_node_count();
    int count = 0;
    int ret = 0;
    int i;

    // Reset the global state
    memset(prealloc_node_enabled, 0, sizeof(prealloc_node_enabled));
    prealloc_nodes_parsed = false;

    // Handle empty parameter: enable all nodes
    if (!prealloc_numa_nodes || !*prealloc_numa_nodes) {
        pr_info("prealloc_numa_nodes is empty, will pre-allocate all nodes\n");

        if (node_count <= 0) {
            pr_err("Invalid NUMA node count: %d\n", node_count);
            return -EINVAL;
        }

        for (i = 0; i < node_count && i < MAX_NUMA_NODES; i++) {
            valid_nodes[i] = true;
        }

        count = node_count;
        goto update_state;
    }

    // Allocate buffer
    buffer = kstrdup(prealloc_numa_nodes, GFP_KERNEL);
    if (!buffer) {
        pr_err("Failed to allocate memory for parsing numa node list\n");
        return -ENOMEM;
    }

    if (node_count <= 0) {
        pr_err("Invalid NUMA node count: %d\n", node_count);
        ret = -EINVAL;
        goto out;
    }

    tmp = buffer;
    while ((token = strsep(&tmp, ",")) != NULL) {
        int node;

        ret = parse_numa_token(token, &node);
        if (ret)
            goto out;

        if (node < 0 || node >= node_count) {
            pr_err("Node %d is out of range (valid: 0-%d)\n", node, node_count - 1);
            ret = -EINVAL;
            goto out;
        }

        if (valid_nodes[node]) {
            pr_warn("Duplicate node %d found, skipping\n", node);
            continue;
        }

        valid_nodes[node] = true;
        count++;
    }

    if (count == 0) {
        pr_err("No valid NUMA nodes specified in '%s'\n", prealloc_numa_nodes);
        ret = -EINVAL;
        goto out;
    }

update_state:
    ret = update_global_state(valid_nodes, node_count);
    if (ret)
        goto out;

    pr_info("Successfully parsed %d valid NUMA nodes\n", count);
    ret = 0;

out:
    kfree(buffer);
    return ret;
}

// - validate_module_params
static int __init validate_module_params(void)
{
    int ret = 0;
    // Validate max_numa_capacity
    if (max_numa_capacity <= 0 ||
        (max_numa_capacity << MB_SHIFT) % MEM_POOL_ALIGNMENT != 0)
    {
        pr_err("max_numa_capacity is invalid: %u\n", max_numa_capacity);
        return -EINVAL;
    }

    // Validate pre_alloc_size
    if (pre_alloc_size > max_numa_capacity ||
        (pre_alloc_size << MB_SHIFT) % MEM_POOL_ALIGNMENT != 0)
    {
        pr_err("pre_alloc_size is invalid: %u\n", pre_alloc_size);
        return -EINVAL;
    }

    // Validate preinit_cache_levels
    if (preinit_cache_levels >= (1 << FIX_CACHE_NUM))
    {
        pr_err("preinit_cache_levels is invalid: 0x%x (valid range: 0x0 - 0x%x)\n",
               preinit_cache_levels, (1 << FIX_CACHE_NUM) - 1);
        return -EINVAL;
    }

    // Validate and parse prealloc_numa_nodes
    ret = parse_and_validate_prealloc_numa_nodes();
    if (ret)
    {
        return ret;
    }

    return 0;
}

// - get_numa_node_count
inline int get_numa_node_count(void)
{
    return atomic_read(&l0_cache_tuner.numa_node_count);
}

// - compare_nodes
static int compare_nodes(const void *a, const void *b)
{
    const int node1 = *(const int *)a;
    const int node2 = *(const int *)b;

    const int dist1 = node_distance(base_node, node1);
    const int dist2 = node_distance(base_node, node2);

    return (dist1 != dist2) ? (dist1 - dist2) : (node1 - node2);
}

// - free_numa_distance_matrix
static void free_numa_distance_matrix(int **matrix, const int node_count)
{
    int i;
    if (!matrix)
    {
        return;
    }

    for (i = 0; i < node_count; i++)
    {
        if (matrix[i])
        {
            kfree(matrix[i]);
            matrix[i] = NULL;
        }
    }
    kfree(matrix);
    matrix = NULL;
}

// - alloc_numa_distance_matrix
static int **alloc_numa_distance_matrix(int node_count)
{
    int **matrix;
    int i;

    if (node_count <= 0 || node_count > MAX_NUMA_NODES) {
        pr_err("NUMA distance matrix allocation: Invalid node count %d\n", node_count);
        return NULL;
    }

    matrix = kcalloc(node_count, sizeof(int *), GFP_KERNEL);
    if (!matrix) {
        pr_err("NUMA distance matrix allocation: Failed to allocate matrix array for %d nodes\n", node_count);
        return NULL;
    }

    for (i = 0; i < node_count; i++)
    {
        matrix[i] = kmalloc_array(node_count, sizeof(int), GFP_KERNEL);
        if (!matrix[i])
        {
            pr_err("NUMA distance matrix allocation: Failed to allocate row %d for %d nodes, freeing previously allocated rows\n", 
                   i, node_count);
            free_numa_distance_matrix(matrix, i);
            return NULL;
        }
    }

    return matrix;
}

// - build_numa_distance_matrix
static int build_numa_distance_matrix(void)
{
    int node_count = get_numa_node_count();
    int **temp_matrix;
    int **old_matrix;
    int i, j, idx;

    temp_matrix = alloc_numa_distance_matrix(node_count);
    if (!temp_matrix)
        return -ENOMEM;
    for_each_online_node(i)
    {
        if (i >= node_count) {
            pr_warn("Skipping node %d, exceeds node count %d\n", i, node_count);
            continue;
        }
        idx = 0;
        temp_matrix[i][idx++] = i;

        for_each_online_node(j)
        {
            if (j != i && idx < node_count) {
                temp_matrix[i][idx++] = j;
            }
        }

        base_node = i;
        if (node_count > 1) {
            sort(temp_matrix[i] + 1, node_count - 1, sizeof(int), compare_nodes, NULL);
        }
    }

    // RCU atomic update
    old_matrix = l0_cache_tuner.numa_distance_matrix;
    rcu_assign_pointer(l0_cache_tuner.numa_distance_matrix, temp_matrix);

    // Delayed free of old matrix
    if (old_matrix) {
        synchronize_rcu();
        free_numa_distance_matrix(old_matrix, atomic_read(&l0_cache_tuner.numa_node_count));
    }

    return 0;
}

// - numa_info_init
static int numa_info_init(void)
{
    int i, ret, node;
    l0_cache_tuner.numa_distance_matrix = NULL;
    ret = build_numa_distance_matrix();
    if (ret != 0)
    {
        return ret;
    }

    // Print sorting results
    pr_debug("NUMA nodes sorted by distance:\n");
    for_each_online_node(node)
    {
        pr_debug("Node %d: [", node);
        for (i = 0; i < get_numa_node_count(); i++)
        {
            pr_cont("%s%d", i ? " " : "", l0_cache_tuner.numa_distance_matrix[node][i]);
        }
        pr_cont("]\n");
    }

    return 0;
}

// - numa_info_destroy
static void numa_info_destroy(void)
{
    int **matrix = l0_cache_tuner.numa_distance_matrix;
    if (matrix)
    {
        RCU_ASSIGN_POINTER(l0_cache_tuner.numa_distance_matrix, NULL);
        SYNCHRONIZE_RCU();
        free_numa_distance_matrix(matrix, atomic_read(&l0_cache_tuner.numa_node_count));
    }
    atomic_set(&l0_cache_tuner.numa_node_count, 0);
}


// - cache_tuner_init
static int cache_tuner_init(void)
{
    int ret = 0;
    struct numa_pool_set *new_pool_list;

    // Allocate memory using normal allocation
    new_pool_list = kcalloc(get_numa_node_count(), sizeof(numa_pool_set), GFP_KERNEL);
    if (!new_pool_list) {
        pr_err("Failed to allocate memory for numa_pool_list\n");
        return -ENOMEM;
    }

    // Initialize the new pool list
    l0_cache_tuner.numa_pool_list = new_pool_list;
    ret = numa_pool_list_init();
    if (ret != 0) {
        pr_err("Initialization numa_pool_list_init failed with error code: %d\n", ret);
        kfree(new_pool_list);
        l0_cache_tuner.numa_pool_list = NULL;
        return ret;
    }
    
    return 0;
}

// - cache_tuner_destroy
static void cache_tuner_destroy(void)
{
    struct numa_pool_set *old_pool_list;

    // Atomically replace with NULL and wait for RCU grace period
    old_pool_list = l0_cache_tuner.numa_pool_list;
    if (old_pool_list) {
        // Use RCU to safely free the old data
        rcu_assign_pointer(l0_cache_tuner.numa_pool_list, NULL);
        synchronize_rcu();                  // Wait for all readers to finish
        numa_pool_list_destroy(old_pool_list); // Now safe to free
    }
}

// - l0_mem_pool_init
static int __init l0_mem_pool_init(void)
{
    int ret;
    
    pr_info("L0 memory pool init: max_capacity=%uMB, pre_alloc=%uMB, preinit_cache=0x%x, prealloc_nodes=\"%s\"\n",
            max_numa_capacity, pre_alloc_size, preinit_cache_levels, prealloc_numa_nodes);
    
    // Initialize NUMA node count first
    ret = early_numa_node_count_init();
    if (ret != 0) {
        return ret;
    }

    ret = validate_module_params();
    if (ret) {
        pr_err("Module parameter validation failed: %d\n", ret);
        return ret;
    }

    ret = numa_info_init();
    if (ret != 0) {
        pr_err("NUMA info initialization failed: %d\n", ret);
        return ret;
    }

    ret = cache_tuner_init();
    if (ret != 0) {
        pr_err("Cache tuner initialization failed: %d\n", ret);
        goto fail_numa_info;
    }
    
    ret = l0_pool_dump_init();
    if (ret != 0) {
        pr_err("Pool dump initialization failed: %d\n", ret);
        goto fail_cache_tuner;
    }

    if (pre_alloc_size > 0) {
        ret = pre_alloc_numa_pool();
        if (ret != 0) {
            pr_err("Pre-allocation failed: %d\n", ret);
            goto fail_dump;
        }
    }
    
    pr_info("L0 memory pool initialization completed successfully\n");
    return 0;

fail_dump:
    // Cleanup dump functionality if needed
fail_cache_tuner:
    cache_tuner_destroy();
fail_numa_info:
    numa_info_destroy();
    return ret;
}

// - l0_mem_pool_exit
static void __exit l0_mem_pool_exit(void)
{
    pr_info("L0 memory pool module exit\n");
    cache_tuner_destroy();
    numa_info_destroy();
    if (dump_l0_pool_kobj) {
        kobject_put(dump_l0_pool_kobj);
    }
    pr_info("L0 memory pool module cleanup completed\n");
}

// - l0_kmalloc
/**
 * Allocates a specified size of memory from the kernel L0 memory pool
 *
 * @param size Size of memory to allocate
 * @param expect_node_id Expected NUMA node ID
 * @return Virtual address of allocated memory, NULL on failure
 */
void *l0_kmalloc(const ssize_t size, const int expect_node_id)
{
    ssize_t req_size;
    phys_addr_t page_pa;
    numa_pool *pool_handle;
    void *vaddr;

    if (unlikely(size <= 0)) {
        pr_err("Invalid size: %zd\n", size);
        return NULL;
    }

    if (unlikely(expect_node_id < 0 || expect_node_id >= get_numa_node_count())) {
        pr_err("Invalid node id: %d\n", expect_node_id);
        return NULL;
    }

    // Align size if necessary
    req_size = (size > MAX_FIX_ALLOCATOR_SIZE) ? 
               ALIGN(size, MEM_POOL_ALIGNMENT) : size;

    // Allocate memory page from L0 mem pool
    page_pa = find_and_alloc_from_numa_pool_list(req_size, &pool_handle, 
                                                 expect_node_id, false);
    if (unlikely(page_pa == 0)) {
        return NULL;
    }
    
    vaddr = phys_to_virt(page_pa);
    return vaddr;
}
EXPORT_SYMBOL_GPL(l0_kmalloc);

// - l0_kfree
/**
 * Frees memory allocated by l0_kmalloc
 *
 * @param mem Memory block to free
 */
void l0_kfree(void *mem)
{
    phys_addr_t base_phys;

    if (unlikely(!mem)) {
        pr_err("l0_kfree: Invalid NULL pointer\n");
        return;
    }

    // Check if memory is valid kernel memory
    if (unlikely(!virt_addr_valid(mem))) {
        pr_err("l0_kfree: Invalid memory address: %p\n", mem);
        return;
    }

    base_phys = virt_to_phys(mem);
    if (unlikely(base_phys == 0)) {
        pr_err("l0_kfree: Failed to convert virtual to physical address: %p\n", mem);
        return;
    }

    find_and_free_to_numa_pool_list(base_phys);
}
EXPORT_SYMBOL_GPL(l0_kfree);

// - l0_alloc_pages
/**
 * Allocates memory pages from L0 memory pool
 *
 * @param page_cnt Number of pages to allocate
 * @param expect_node_id Expected NUMA node ID
 * @return Pointer to allocated page structure, NULL on failure
 */
struct page *l0_alloc_pages(const int page_cnt, const int expect_node_id)
{
    ssize_t req_size;
    phys_addr_t page_pa;
    numa_pool *pool_handle;
    struct page *page;

    if (unlikely(page_cnt <= 0)) {
        pr_err("Invalid page count: %d\n", page_cnt);
        return NULL;
    }

    if (unlikely(expect_node_id < 0 || expect_node_id >= get_numa_node_count())) {
        pr_err("Invalid node id: %d\n", expect_node_id);
        return NULL;
    }

    req_size = (ssize_t)page_cnt << PAGE_SHIFT;

    // Allocate memory page from L0 mem pool
    page_pa = find_and_alloc_from_numa_pool_list(req_size, &pool_handle, 
                                                 expect_node_id, true);
    if (unlikely(page_pa == 0)) {
        return NULL;
    }

    page = phys_to_page(page_pa);
    if (unlikely(!page)) {
        pr_err("Failed to convert physical address to page\n");
        return NULL;
    }

    return page;
}
EXPORT_SYMBOL_GPL(l0_alloc_pages);

// - l0_free_pages
/**
 * Frees pages allocated by l0_alloc_pages
 *
 * @param page Page structure to free
 */
void l0_free_pages(struct page *page)
{
    phys_addr_t base_phys;

    // Validate input parameter
    if (unlikely(!page)) {
        pr_err("l0_free_pages: Invalid NULL page parameter\n");
        return;
    }
    
    base_phys = page_to_phys(page);
    if (unlikely(base_phys == 0)) {
        pr_err("l0_free_pages: Failed to convert page to physical address\n");
        return;
    }
    
    find_and_free_to_numa_pool_list(base_phys);
}
EXPORT_SYMBOL_GPL(l0_free_pages);


module_init(l0_mem_pool_init);
module_exit(l0_mem_pool_exit);

MODULE_LICENSE("GPL v2");
MODULE_AUTHOR("YunYi Yang yangyunyi2@huawei.com");
MODULE_DESCRIPTION("HiSilicon SoC L0 mem pool driver");